package com.yfp.client.config;

import com.yfp.client.domain.dto.SerializableMessageDTO;
import jakarta.annotation.Resource;
import org.redisson.api.RMap;
import org.redisson.api.RedissonClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.stereotype.Component;

import java.util.Collections;
import java.util.List;


@Component
public class AIRedisChatMemory implements ChatMemory {

    @Resource
    private RedissonClient redissonClient;

    @Override
    public void add(String conversationId, List<Message> messages) {
        RMap<String, List<SerializableMessageDTO>> aiRedisChatMemory = redissonClient.getMap("AIRedisChatMemory");
        List<SerializableMessageDTO> serializableMessages = messages.stream().map(SerializableMessageDTO::new).toList();
        aiRedisChatMemory.put(conversationId, serializableMessages);
    }

    @Override
    public void add(String conversationId, Message message) {
        RMap<String, List<SerializableMessageDTO>> aiRedisChatMemory = redissonClient.getMap("AIRedisChatMemory");
        SerializableMessageDTO serializableMessageDTO = new SerializableMessageDTO();

        aiRedisChatMemory.put(conversationId, Collections.singletonList(serializableMessageDTO));
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        RMap<String, List<SerializableMessageDTO>> aiRedisChatMemory = redissonClient.getMap("AIRedisChatMemory");
        List<SerializableMessageDTO> all = aiRedisChatMemory.get(conversationId);
        return all != null ? all.stream().map(SerializableMessageDTO::getMessage).skip(Math.max(0, all.size() - lastN)).toList() : List.of();
    }

    @Override
    public void clear(String conversationId) {
        RMap<String, List<SerializableMessageDTO>> aiRedisChatMemory = redissonClient.getMap("AIRedisChatMemory");
        aiRedisChatMemory.remove(conversationId);
    }

}
